#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从gitee下载分包文件，合并并安装
"""

import requests
import json
import os
import subprocess
from pathlib import Path
import time

# 修改为raw链接格式
BASE_URL = "https://gitee.com/qiufengruge/comfyui_ctrlscripts/raw/master/split_exes"
# 获取当前脚本所在目录，然后上两级到项目根目录，再上一级到上级目录
SCRIPT_DIR = Path(__file__).parent
PROJECT_ROOT = SCRIPT_DIR.parent.parent
PARENT_ROOT = PROJECT_ROOT.parent  # 上三级目录
DOWNLOAD_DIR = PARENT_ROOT / "split_exes"
TOOLS_DIR = PARENT_ROOT / "tools"  # 合并后的文件也放到上三级目录

def download_file(url, local_path):
    """下载单个文件"""
    try:
        print(f"下载: {url}")
        r = requests.get(url, stream=True)
        r.raise_for_status()
        local_path.parent.mkdir(parents=True, exist_ok=True)
        with open(local_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"✅ 下载成功: {local_path}")
        return True
    except Exception as e:
        print(f"❌ 下载失败: {url}，原因: {e}")
        return False

def download_config():
    """下载config.json"""
    config_url = f"{BASE_URL}/config.json"
    local_path = DOWNLOAD_DIR / "config.json"
    return download_file(config_url, local_path)

def download_all_chunks_from_config():
    """从config.json读取并下载所有分包文件"""
    config_path = DOWNLOAD_DIR / "config.json"
    if not config_path.exists():
        print("❌ config.json 未找到")
        return False

    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    files_to_download = set()
    # 遍历所有文件
    for file_info in config["files"]:
        original = file_info["original"]
        split_dir = Path(original).stem
        files_to_download.add(f"{split_dir}/info.json")
        for chunk in file_info["chunks"]:
            files_to_download.add(f"{split_dir}/{chunk['file']}")

    print(f"需要下载 {len(files_to_download)} 个文件")
    
    # 下载所有分包
    success_count = 0
    for rel_path in files_to_download:
        url = f"{BASE_URL}/{rel_path}"
        local_path = DOWNLOAD_DIR / rel_path
        if download_file(url, local_path):
            success_count += 1
        time.sleep(0.1)  # 避免请求过快
    
    print(f"下载完成: {success_count}/{len(files_to_download)}")
    return success_count == len(files_to_download)

def merge_files():
    """合并分包文件"""
    print("\n开始合并分包文件...")
    config_file = DOWNLOAD_DIR / "config.json"
    if not config_file.exists():
        print("❌ 配置文件不存在")
        return False

    with open(config_file, 'r', encoding='utf-8') as f:
        config = json.load(f)

    TOOLS_DIR.mkdir(exist_ok=True)

    for file_info in config['files']:
        original_name = file_info['original']
        split_dir = DOWNLOAD_DIR / Path(original_name).stem
        print(f"合并: {original_name}")

        info_file = split_dir / "info.json"
        if not info_file.exists():
            print(f"❌ 元数据文件不存在: {info_file}")
            continue

        with open(info_file, 'r', encoding='utf-8') as f:
            info = json.load(f)

        output_file = TOOLS_DIR / original_name
        with open(output_file, 'wb') as out:
            for chunk_info in info['chunks']:
                chunk_file = split_dir / chunk_info['file']
                if chunk_file.exists():
                    with open(chunk_file, 'rb') as cf:
                        out.write(cf.read())
                    print(f"  合并分包: {chunk_info['file']}")
                else:
                    print(f"❌ 分包文件不存在: {chunk_file}")
                    return False
        print(f"✅ 合并完成: {output_file}")
    
    print("✅ 所有文件合并完成!")
    
    # 合并完成后删除下载的分包文件
    print("\n清理下载的分包文件...")
    try:
        import shutil
        if DOWNLOAD_DIR.exists():
            shutil.rmtree(DOWNLOAD_DIR)
            print(f"✅ 已删除下载目录: {DOWNLOAD_DIR}")
        else:
            print(f"⚠️  下载目录不存在: {DOWNLOAD_DIR}")
    except Exception as e:
        print(f"❌ 删除下载文件失败: {e}")
    
    return True

def install_exe_files():
    """安装exe文件"""
    print("\n开始安装exe文件...")
    if not TOOLS_DIR.exists():
        print("❌ tools目录不存在")
        return False

    exe_files = list(TOOLS_DIR.glob("*.exe"))
    if not exe_files:
        print("❌ 未找到exe文件")
        return False

    print(f"找到 {len(exe_files)} 个exe文件:")
    for exe_file in exe_files:
        print(f"  - {exe_file.name}")

    print("\n开始安装...")
    for exe_file in exe_files:
        print(f"\n正在安装: {exe_file.name}")
        try:
            print("尝试静默安装...")
            result = subprocess.run([
                str(exe_file),
                "/S",
                f"/D=C:\\Program Files\\{exe_file.stem}"
            ], capture_output=True, text=True, timeout=60)
            if result.returncode == 0:
                print(f"✅ {exe_file.name} 静默安装完成")
            else:
                print("静默安装失败，尝试普通安装...")
                subprocess.run([str(exe_file)], timeout=300)
                print(f"✅ {exe_file.name} 普通安装完成")
        except subprocess.TimeoutExpired:
            print(f"⚠️  {exe_file.name} 安装超时")
        except Exception as e:
            print(f"❌ {exe_file.name} 安装失败: {e}")
    print("\n✅ 安装完成!")
    return True

def main():
    print("=" * 60)
    print("        Gitee分包文件下载和安装工具")
    print("=" * 60)
    print(f"源地址: {BASE_URL}")

    # 1. 下载config.json
    print("\n步骤1: 下载config.json")
    if not download_config():
        print("❌ config.json 下载失败")
        return

    # 2. 下载所有分包
    print("\n步骤2: 下载所有分包文件")
    if not download_all_chunks_from_config():
        print("❌ 分包文件下载失败")
        return

    # 3. 合并
    print("\n步骤3: 合并分包文件")
    if not merge_files():
        print("❌ 合并失败")
        return

    # 4. 安装
    print("\n步骤4: 安装exe文件")
    if not install_exe_files():
        print("❌ 安装失败")
        return

    print("\n" + "=" * 60)
    print("✅ 所有操作完成!")
    print("=" * 60)

if __name__ == "__main__":
    main() 